#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
LinEnum Integration Module
LinEnum集成模块
"""

import os
import re
import json
import time
import subprocess
from typing import Dict, List, Any, Optional
from pathlib import Path
from .tool_manager import ExternalToolManager

class LinEnumIntegration:
    """LinEnum集成类"""
    
    def __init__(self, tool_manager: ExternalToolManager = None):
        self.tool_manager = tool_manager or ExternalToolManager()
        self.tool_name = "linenum"
        self.last_scan_results = None
        self.scan_history = []
    
    def is_available(self) -> bool:
        """检查LinEnum是否可用"""
        status = self.tool_manager.get_tool_status(self.tool_name)
        return status == "ready"
    
    def ensure_available(self) -> bool:
        """确保LinEnum可用"""
        if self.is_available():
            return True
        
        print("📥 LinEnum不可用，正在下载...")
        return self.tool_manager.download_tool(self.tool_name)
    
    def run_standard_enum(self, output_file: str = None) -> Dict[str, Any]:
        """运行标准枚举"""
        if not self.ensure_available():
            return {
                "success": False,
                "error": "LinEnum不可用",
                "results": {}
            }
        
        args = []  # 标准参数
        
        if output_file:
            # LinEnum使用重定向输出
            args = [">", output_file]
        
        return self._execute_scan(args, "standard")
    
    def run_thorough_enum(self, output_file: str = None) -> Dict[str, Any]:
        """运行详细枚举"""
        if not self.ensure_available():
            return {
                "success": False,
                "error": "LinEnum不可用",
                "results": {}
            }
        
        args = ["-t"]  # 详细模式参数
        
        if output_file:
            args.extend([">", output_file])
        
        return self._execute_scan(args, "thorough")
    
    def run_keyword_search(self, keywords: List[str], output_file: str = None) -> Dict[str, Any]:
        """运行关键词搜索"""
        if not self.ensure_available():
            return {
                "success": False,
                "error": "LinEnum不可用",
                "results": {}
            }
        
        # LinEnum使用-k参数进行关键词搜索
        args = ["-k", ",".join(keywords)]
        
        if output_file:
            args.extend([">", output_file])
        
        return self._execute_scan(args, "keyword_search")
    
    def run_export_enum(self, export_dir: str = "/tmp/linenum_export") -> Dict[str, Any]:
        """运行导出枚举"""
        if not self.ensure_available():
            return {
                "success": False,
                "error": "LinEnum不可用",
                "results": {}
            }
        
        # 创建导出目录
        os.makedirs(export_dir, exist_ok=True)
        
        args = ["-e", export_dir]  # 导出模式参数
        
        return self._execute_scan(args, "export")
    
    def run_custom_enum(self, custom_args: List[str], scan_type: str = "custom") -> Dict[str, Any]:
        """运行自定义枚举"""
        if not self.ensure_available():
            return {
                "success": False,
                "error": "LinEnum不可用",
                "results": {}
            }
        
        return self._execute_scan(custom_args, scan_type)
    
    def _execute_scan(self, args: List[str], scan_type: str) -> Dict[str, Any]:
        """执行扫描"""
        print(f"🔍 开始LinEnum {scan_type}枚举...")
        start_time = time.time()
        
        # 执行LinEnum
        result = self.tool_manager.execute_tool(
            self.tool_name, 
            args, 
            timeout=300,  # 5分钟超时
            capture_output=True
        )
        
        execution_time = time.time() - start_time
        
        if result["success"]:
            # 解析结果
            parsed_results = self._parse_output(result["stdout"])
            
            scan_result = {
                "success": True,
                "scan_type": scan_type,
                "execution_time": execution_time,
                "raw_output": result["stdout"],
                "parsed_results": parsed_results,
                "timestamp": time.time(),
                "command": result["command"]
            }
            
            # 保存结果
            self.last_scan_results = scan_result
            self.scan_history.append(scan_result)
            
            print(f"✅ LinEnum {scan_type}枚举完成 (耗时: {execution_time:.2f}s)")
            print(f"📊 发现 {len(parsed_results.get('findings', []))} 个项目")
            
            return scan_result
        else:
            error_result = {
                "success": False,
                "error": result["error"],
                "scan_type": scan_type,
                "execution_time": execution_time,
                "stderr": result["stderr"],
                "returncode": result["returncode"]
            }
            
            print(f"❌ LinEnum {scan_type}枚举失败: {result['error']}")
            
            return error_result
    
    def _parse_output(self, output: str) -> Dict[str, Any]:
        """解析LinEnum输出"""
        parsed = {
            "system_info": {},
            "kernel_info": {},
            "user_info": {},
            "environment_info": {},
            "job_info": {},
            "network_info": {},
            "service_info": {},
            "software_info": {},
            "file_info": {},
            "development_tools": {},
            "findings": [],
            "interesting_files": []
        }
        
        try:
            lines = output.split('\n')
            current_section = None
            
            for line in lines:
                line = line.strip()
                
                # 检测章节标题
                if self._is_section_header(line):
                    current_section = self._extract_section_name(line)
                    continue
                
                # 根据当前章节解析内容
                if current_section:
                    self._parse_section_content(line, current_section, parsed)
                
                # 检测有趣的发现
                finding = self._extract_finding(line)
                if finding:
                    parsed["findings"].append(finding)
                
                # 检测有趣的文件
                interesting_file = self._extract_interesting_file(line)
                if interesting_file:
                    parsed["interesting_files"].append(interesting_file)
            
            # 统计信息
            parsed["statistics"] = {
                "total_findings": len(parsed["findings"]),
                "interesting_files_count": len(parsed["interesting_files"]),
                "sections_parsed": len([k for k, v in parsed.items() if isinstance(v, dict) and v])
            }
            
        except Exception as e:
            print(f"⚠️ 输出解析警告: {e}")
            parsed["parse_error"] = str(e)
        
        return parsed
    
    def _is_section_header(self, line: str) -> bool:
        """检测是否为章节标题"""
        section_patterns = [
            r"#{10,}",  # 井号分隔符
            r"\[-\]",   # [-] 标记
            r"\[!\]",   # [!] 标记
            r"\[\+\]",  # [+] 标记
            r"={10,}",  # 等号分隔符
        ]
        
        for pattern in section_patterns:
            if re.search(pattern, line):
                return True
        
        return False
    
    def _extract_section_name(self, line: str) -> str:
        """提取章节名称"""
        # 移除特殊字符，提取核心内容
        clean_line = re.sub(r'[#\[\]\-\!\+\=]+', '', line).strip()
        
        # 常见章节映射
        section_mapping = {
            "system information": "system_info",
            "kernel information": "kernel_info",
            "user information": "user_info",
            "environmental": "environment_info",
            "jobs": "job_info",
            "networking": "network_info",
            "services": "service_info",
            "software": "software_info",
            "interesting files": "file_info",
            "development tools": "development_tools"
        }
        
        for key, value in section_mapping.items():
            if key in clean_line.lower():
                return value
        
        return "general"
    
    def _parse_section_content(self, line: str, section: str, parsed: Dict[str, Any]):
        """解析章节内容"""
        if not line:
            return
        
        try:
            if section == "system_info":
                self._parse_system_info(line, parsed["system_info"])
            elif section == "kernel_info":
                self._parse_kernel_info(line, parsed["kernel_info"])
            elif section == "user_info":
                self._parse_user_info(line, parsed["user_info"])
            elif section == "environment_info":
                self._parse_environment_info(line, parsed["environment_info"])
            elif section == "job_info":
                self._parse_job_info(line, parsed["job_info"])
            elif section == "network_info":
                self._parse_network_info(line, parsed["network_info"])
            elif section == "service_info":
                self._parse_service_info(line, parsed["service_info"])
            elif section == "software_info":
                self._parse_software_info(line, parsed["software_info"])
            elif section == "file_info":
                self._parse_file_info(line, parsed["file_info"])
            elif section == "development_tools":
                self._parse_development_tools(line, parsed["development_tools"])
        except Exception:
            pass  # 忽略解析错误
    
    def _parse_system_info(self, line: str, system_info: Dict[str, Any]):
        """解析系统信息"""
        if "hostname" in line.lower():
            system_info["hostname"] = self._extract_value(line)
        elif "operating system" in line.lower():
            system_info["os"] = self._extract_value(line)
        elif "architecture" in line.lower():
            system_info["architecture"] = self._extract_value(line)
        elif "uptime" in line.lower():
            system_info["uptime"] = self._extract_value(line)
    
    def _parse_kernel_info(self, line: str, kernel_info: Dict[str, Any]):
        """解析内核信息"""
        if "kernel version" in line.lower():
            kernel_info["version"] = self._extract_value(line)
        elif "kernel release" in line.lower():
            kernel_info["release"] = self._extract_value(line)
    
    def _parse_user_info(self, line: str, user_info: Dict[str, Any]):
        """解析用户信息"""
        if "current user" in line.lower():
            user_info["current_user"] = self._extract_value(line)
        elif "current user id" in line.lower():
            user_info["current_uid"] = self._extract_value(line)
        elif "users logged on" in line.lower():
            if "logged_users" not in user_info:
                user_info["logged_users"] = []
            user_info["logged_users"].append(line.strip())
        elif "user accounts" in line.lower():
            if "all_users" not in user_info:
                user_info["all_users"] = []
            user_info["all_users"].append(line.strip())
    
    def _parse_environment_info(self, line: str, env_info: Dict[str, Any]):
        """解析环境信息"""
        if "path" in line.lower() and "=" in line:
            env_info["path"] = self._extract_value(line)
        elif "shell" in line.lower() and "=" in line:
            env_info["shell"] = self._extract_value(line)
        elif "=" in line and len(line.split("=")) == 2:
            key, value = line.split("=", 1)
            env_info[key.strip().lower()] = value.strip()
    
    def _parse_job_info(self, line: str, job_info: Dict[str, Any]):
        """解析任务信息"""
        if "cron" in line.lower():
            if "cron_jobs" not in job_info:
                job_info["cron_jobs"] = []
            job_info["cron_jobs"].append(line.strip())
        elif "at" in line.lower() and "job" in line.lower():
            if "at_jobs" not in job_info:
                job_info["at_jobs"] = []
            job_info["at_jobs"].append(line.strip())
    
    def _parse_network_info(self, line: str, network_info: Dict[str, Any]):
        """解析网络信息"""
        if "interface" in line.lower():
            if "interfaces" not in network_info:
                network_info["interfaces"] = []
            network_info["interfaces"].append(line.strip())
        elif "listening" in line.lower():
            if "listening_services" not in network_info:
                network_info["listening_services"] = []
            network_info["listening_services"].append(line.strip())
        elif "route" in line.lower():
            if "routes" not in network_info:
                network_info["routes"] = []
            network_info["routes"].append(line.strip())
    
    def _parse_service_info(self, line: str, service_info: Dict[str, Any]):
        """解析服务信息"""
        if "running" in line.lower() and "service" in line.lower():
            if "running_services" not in service_info:
                service_info["running_services"] = []
            service_info["running_services"].append(line.strip())
        elif "systemd" in line.lower():
            if "systemd_services" not in service_info:
                service_info["systemd_services"] = []
            service_info["systemd_services"].append(line.strip())
    
    def _parse_software_info(self, line: str, software_info: Dict[str, Any]):
        """解析软件信息"""
        if "installed" in line.lower():
            if "installed_packages" not in software_info:
                software_info["installed_packages"] = []
            software_info["installed_packages"].append(line.strip())
        elif "version" in line.lower():
            if "software_versions" not in software_info:
                software_info["software_versions"] = []
            software_info["software_versions"].append(line.strip())
    
    def _parse_file_info(self, line: str, file_info: Dict[str, Any]):
        """解析文件信息"""
        if "writable" in line.lower():
            if "writable_files" not in file_info:
                file_info["writable_files"] = []
            file_info["writable_files"].append(line.strip())
        elif "suid" in line.lower():
            if "suid_files" not in file_info:
                file_info["suid_files"] = []
            file_info["suid_files"].append(line.strip())
        elif "config" in line.lower():
            if "config_files" not in file_info:
                file_info["config_files"] = []
            file_info["config_files"].append(line.strip())
    
    def _parse_development_tools(self, line: str, dev_tools: Dict[str, Any]):
        """解析开发工具"""
        tools = ["gcc", "python", "perl", "ruby", "java", "node", "php"]
        
        for tool in tools:
            if tool in line.lower():
                if "available_tools" not in dev_tools:
                    dev_tools["available_tools"] = []
                dev_tools["available_tools"].append(line.strip())
                break
    
    def _extract_value(self, line: str) -> str:
        """提取值"""
        if ":" in line:
            return line.split(":", 1)[1].strip()
        elif "=" in line:
            return line.split("=", 1)[1].strip()
        else:
            return line.strip()
    
    def _extract_finding(self, line: str) -> Optional[Dict[str, Any]]:
        """提取发现"""
        # 检测重要发现
        important_patterns = [
            (r"password.*file", "high", "credentials"),
            (r"private.*key", "high", "credentials"),
            (r"suid.*root", "high", "privilege_escalation"),
            (r"sudo.*nopasswd", "high", "privilege_escalation"),
            (r"writable.*etc", "medium", "file_permissions"),
            (r"world.*writable", "medium", "file_permissions"),
            (r"interesting.*file", "low", "information_disclosure")
        ]
        
        for pattern, priority, category in important_patterns:
            if re.search(pattern, line, re.IGNORECASE):
                return {
                    "type": "enumeration_finding",
                    "priority": priority,
                    "category": category,
                    "description": line.strip(),
                    "line": line.strip()
                }
        
        return None
    
    def _extract_interesting_file(self, line: str) -> Optional[str]:
        """提取有趣的文件"""
        interesting_patterns = [
            r"/etc/passwd",
            r"/etc/shadow",
            r"/etc/sudoers",
            r"\.ssh/",
            r"\.bash_history",
            r"\.mysql_history",
            r"config\.php",
            r"wp-config\.php"
        ]
        
        for pattern in interesting_patterns:
            if re.search(pattern, line, re.IGNORECASE):
                return line.strip()
        
        return None
    
    def get_last_results(self) -> Optional[Dict[str, Any]]:
        """获取最后一次扫描结果"""
        return self.last_scan_results
    
    def get_scan_history(self) -> List[Dict[str, Any]]:
        """获取扫描历史"""
        return self.scan_history
    
    def export_results(self, file_path: str, format: str = "json"):
        """导出结果"""
        if not self.last_scan_results:
            print("❌ 没有可导出的结果")
            return False
        
        try:
            if format.lower() == "json":
                with open(file_path, 'w', encoding='utf-8') as f:
                    json.dump(self.last_scan_results, f, indent=2, ensure_ascii=False)
            
            elif format.lower() == "txt":
                with open(file_path, 'w', encoding='utf-8') as f:
                    f.write("LinEnum枚举结果\n")
                    f.write("=" * 50 + "\n\n")
                    f.write(f"枚举类型: {self.last_scan_results['scan_type']}\n")
                    f.write(f"执行时间: {self.last_scan_results['execution_time']:.2f}秒\n")
                    f.write(f"时间戳: {time.ctime(self.last_scan_results['timestamp'])}\n\n")
                    f.write("原始输出:\n")
                    f.write("-" * 30 + "\n")
                    f.write(self.last_scan_results['raw_output'])
            
            print(f"📄 结果已导出到: {file_path}")
            return True
            
        except Exception as e:
            print(f"❌ 导出失败: {e}")
            return False
    
    def get_summary(self) -> Dict[str, Any]:
        """获取枚举摘要"""
        if not self.last_scan_results:
            return {"error": "没有可用的枚举结果"}
        
        parsed = self.last_scan_results.get("parsed_results", {})
        stats = parsed.get("statistics", {})
        
        return {
            "scan_type": self.last_scan_results["scan_type"],
            "execution_time": self.last_scan_results["execution_time"],
            "total_findings": stats.get("total_findings", 0),
            "interesting_files": stats.get("interesting_files_count", 0),
            "sections_parsed": stats.get("sections_parsed", 0),
            "timestamp": self.last_scan_results["timestamp"]
        }
    
    def search_in_results(self, keyword: str) -> List[str]:
        """在结果中搜索关键词"""
        if not self.last_scan_results:
            return []
        
        matches = []
        raw_output = self.last_scan_results.get("raw_output", "")
        
        for line in raw_output.split('\n'):
            if keyword.lower() in line.lower():
                matches.append(line.strip())
        
        return matches

# 使用示例
def example_usage():
    """使用示例"""
    # 创建LinEnum集成实例
    linenum = LinEnumIntegration()
    
    # 检查可用性
    if linenum.is_available():
        print("✅ LinEnum已就绪")
    else:
        print("📥 正在准备LinEnum...")
        if not linenum.ensure_available():
            print("❌ LinEnum准备失败")
            return
    
    # 运行标准枚举
    print("🔍 开始标准枚举...")
    result = linenum.run_standard_enum()
    
    if result["success"]:
        print("✅ 枚举完成")
        
        # 获取摘要
        summary = linenum.get_summary()
        print(f"📊 枚举摘要:")
        print(f"  - 发现总数: {summary['total_findings']}")
        print(f"  - 有趣文件: {summary['interesting_files']}")
        print(f"  - 解析章节: {summary['sections_parsed']}")
        
        # 搜索特定内容
        password_matches = linenum.search_in_results("password")
        if password_matches:
            print(f"🔍 找到 {len(password_matches)} 个密码相关项目")
        
        # 导出结果
        linenum.export_results("linenum_results.json", "json")
        
    else:
        print(f"❌ 枚举失败: {result['error']}")

if __name__ == "__main__":
    example_usage()